#!/usr/bin/env python
# -*- coding: utf_8 -*-

from __future__ import print_function
import rospy
import socket
import math
import pickle
import numpy as np
from std_msgs.msg import Int8, Bool
from robot_jaka.msg import *
from robot_jaka.srv import *
# from g3p_msgs.msg import Gripper

import time

class RobotVisionWrapper:
    def __init__(self):
        rospy.loginfo("==========>START Robot Vision Wrapper, press [Ctrl+C]: exit<==========")
        rospy.on_shutdown(self.cleanUp)

        self.visionSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.visionSocket.settimeout(5)
        self.visionIsConnected = False
        self.isEyeInHand = False

        self.inputResponding = False
        self.cancelResponding = False
        self.objectCompensationXYZ = None
        self.useInputRPY = None
        self.objectRotateRPY = None
        self.objectPrePoseOffsetXYZRPY = None
        self.autoPlace = None
        self.pickPrePose1Joint=None
        self.pickPrePose2Joint=None
        self.pickLatePose1Joint=None
        self.pickLatePose2Joint=None
        self.placeLatePose1Joint=None
        self.placeLatePose2Joint=None

        rospy.wait_for_service("robot_service/robot_end_position")
        rospy.wait_for_service("robot_service/robot_end_move")
        self.handle_position = rospy.ServiceProxy("robot_service/robot_end_position",RobotEndPosition)
        self.handle_move = rospy.ServiceProxy("robot_service/robot_end_move",RobotEndMove)
        self.wrapperInputSub = rospy.Subscriber("/vision_jaka/input",Int8,self.inputCallback,queue_size=1)
        self.robotGripperPub = rospy.Publisher("robot_topic/robot_gripper_close",Bool,queue_size=1)
        # self.robotGripper3Pub = rospy.Publisher("/g3p_msgs/gripper",Gripper,queue_size=5)
    def cleanUp(self):
        rospy.loginfo("==========>Robot Vision Wrapper STOPPED!<==========")
        # rospy.sleep(5)

    def initParam(self, isEyeInHand, objectCompensationXYZ, useInputRPY, objectRotateRPY, objectPrePoseOffsetXYZRPY, autoPlace,
                    pickPrePose1Joint, pickPrePose2Joint, pickLatePose1Joint, pickLatePose2Joint):
        self.isEyeInHand = isEyeInHand
        self.objectCompensationXYZ = tuple(map(float,objectCompensationXYZ.split(",")))
        self.useInputRPY = useInputRPY
        self.objectRotateRPY = tuple(map(float,objectRotateRPY.split(",")))
        self.objectPrePoseOffsetXYZRPY = tuple(map(float,objectPrePoseOffsetXYZRPY.split(",")))
        self.autoPlace = autoPlace
        self.pickPrePose1Joint = tuple(map(float,pickPrePose1Joint.split(",")))
        self.pickPrePose2Joint = tuple(map(float,pickPrePose2Joint.split(",")))
        self.pickLatePose1Joint = tuple(map(float,pickLatePose1Joint.split(",")))
        self.pickLatePose2Joint = tuple(map(float,pickLatePose2Joint.split(",")))
        self.placeLatePose1Joint = tuple(map(float,placeLatePose1Joint.split(",")))
        self.placeLatePose2Joint = tuple(map(float,placeLatePose2Joint.split(",")))

    def inputCallback(self, msg):
        # self.wrapperInputSub.unregister()
        # self.wrapperInputSub = rospy.Subscriber("/vision_rocr6/input",Int8,self.inputCallback,queue_size=1)

        if msg.data == 0:
            self.cancelResponding = True
            self.inputResponding = False
            self.handle_move((-145, -60, 430, 175, 25, 53),10)
            self.cancelResponding = False
            return
            
        if not self.inputResponding:
            self.inputResponding = True
            if msg.data == 99:
                self.robotPlace()
            else:
                if self.visionIsConnected:
                    matrix = self.visionObjectRequest(str(msg.data))
                    if np.all(matrix == 0):
                        rospy.logwarn("The target object is not detected !")
                    elif np.all(matrix == -1):
                        rospy.logwarn("Illegal input, please re-enter!")
                    else:
                        if self.isEyeInHand:
                            xPose = self.handle_position().pos
                            robotCurrentPoseMatrix = self.robotPose2Matrix(xPose)
                            matrix = robotCurrentPoseMatrix.dot(matrix)
                        self.robotPick(matrix)
                        if self.autoPlace:
                            self.robotPlace()
                else:
                    rospy.logfatal("Need to connect the VisionServer first, please try to restart this node !")
            self.inputResponding = False
        else:
            rospy.logwarn("Input is responding, please wait until the response is over !")

    def connectVision(self, visionIPAddress, visionPort):
        rospy.loginfo("Connect to VisionServer ......")
        while not self.visionIsConnected and not rospy.is_shutdown():
            try:
                self.visionSocket.connect((visionIPAddress,visionPort))
                self.visionIsConnected = True
                rospy.loginfo("Connect to VisionServer Success !")
            except Exception as e:
                rospy.logwarn("Connect to VisionServer Failed! %s",e)
            rospy.sleep(5)

    def disconnectVision(self):
        while self.visionIsConnected and not rospy.is_shutdown():
            self.visionSocket.close()
            self.visionIsConnected = False
            rospy.loginfo("Disconnect to VisionServer Success !")
            rospy.sleep(5)

    def visionAction(self,commandString):
        self.visionSocket.sendall(commandString.encode("ascii"))
        recvData=self.visionSocket.recv(1024).decode("ascii")
        rospy.loginfo("tx:[%s]    rx:[%s]",commandString,recvData)
        return recvData

    def startVision(self):
        rospy.loginfo(self.visionAction("start"))

    def stopVision(self):
        rospy.loginfo(self.visionAction("stop"))

    def visionObjectRequest(self,objectNo):
        rospy.loginfo("objectNo:[%s]", objectNo)
        self.visionSocket.sendall(objectNo.encode("ascii"))
        # test=pickle.dumps(matrix,protocol=0)
        rospy.sleep(0.2)
        matrix = pickle.loads(self.visionSocket.recv(1024))
        return matrix

    def robotPose2Matrix(self, pose):
        R=(pose[3]*math.pi/180, pose[4]*math.pi/180, pose[5]*math.pi/180)
        Mx=np.array([[1,                0,					0],
                    [0,		            math.cos(R[0]),	    -math.sin(R[0])],
                    [0,		            math.sin(R[0]),	    math.cos(R[0])]])

        My=np.array([[math.cos(R[1]),   0,		            math.sin(R[1])],
                    [0,                 1,	                0],
                    [-math.sin(R[1]),	0,	                math.cos(R[1])]])

        Mz=np.array([[math.cos(R[2]),   -math.sin(R[2]),	0],
                    [math.sin(R[2]),    math.cos(R[2]),     0],
                    [0,		            0,	                1]])

        M=Mz.dot(My.dot(Mx))
        M=np.array([[M[0,0],    M[0,1], M[0,2], pose[0]*0.001],
                    [M[1,0],    M[1,1], M[1,2], pose[1]*0.001],
                    [M[2,0],    M[2,1], M[2,2], pose[2]*0.001],
                    [0,         0,      0,      1]])
        return M

    def robotMatrix2Pose(self, matrix):
        Rr, Rp, Ry = None, None, None
        if self.useInputRPY:
            Rr = self.objectRotateRPY[0]
            Rp = self.objectRotateRPY[1]
            Ry = self.objectRotateRPY[2]
        else:
            sy = math.sqrt(math.pow(matrix[0,0],2)+math.pow(matrix[1,0],2))
            if not sy<1e-6:
                Rr = math.atan2(matrix[2,1],matrix[2,2])
                Rp = math.atan2(-matrix[2,0],sy)
                Ry = math.atan2(matrix[1,0],matrix[0,0])
            else:
                Rr = math.atan2(-matrix[1,2],matrix[1,1])
                Rp = math.atan2(-matrix[2,0],sy)
                Ry = 0
        return (matrix[0,3]*1000 + self.objectCompensationXYZ[0],
                matrix[1,3]*1000 + self.objectCompensationXYZ[1],
                matrix[2,3]*1000 + self.objectCompensationXYZ[2],
                Rr*180/math.pi, Rp*180/math.pi, Ry*180/math.pi)

    def robotMatrix2PrePose(self, pose):
        return (pose[0]+self.objectPrePoseOffsetXYZRPY[0], pose[1]+self.objectPrePoseOffsetXYZRPY[1],
                pose[2]+self.objectPrePoseOffsetXYZRPY[2], pose[3]+self.objectPrePoseOffsetXYZRPY[3],
                pose[4]+self.objectPrePoseOffsetXYZRPY[4], pose[5]+self.objectPrePoseOffsetXYZRPY[5])

    def robotMove(self,pose):
        if pose != (-1.0,-1.0,-1.0,-1.0,-1.0,-1.0):
            self.handle_move(pose,30)

    def robotGripperChangeState(self, state):
        self.robotGripperPub.publish(state)
        time.sleep(2)

    # def robotGripper3ChangeState(self, state):
    #     gripper3State = Gripper()
    #     gripper3State.state = state
    #     self.robotGripper3Pub.publish(gripper3State)
    #     rospy.loginfo("gripper3ChangeState: %f",gripper3State.state)

    # def robotGripper3ChangeStateL(self, state):
    #     while self.robotStatus.status != 3 and not rospy.is_shutdown():
    #         rospy.sleep(1)
    #     if self.robotStatus.status == 3 and not self.cancelResponding:
    #         self.robotGripper3ChangeState(state)
    #         rospy.sleep(3)

    def robotPick(self, matrix):
        pose = self.robotMatrix2Pose(matrix)
        self.robotGripperChangeState(False)
        self.robotMove(self.pickPrePose1Joint)
        self.robotMove(self.pickPrePose2Joint)
        self.robotMove(self.robotMatrix2PrePose(pose))
        time.sleep(1)
        self.robotMove(pose)
        self.robotGripperChangeState(True)
        self.robotMove(self.robotMatrix2PrePose(pose))

    def robotPlace(self):
        self.robotMove(self.pickLatePose1Joint)
        self.robotMove(self.pickLatePose2Joint)
        self.robotGripperChangeState(False)
        self.robotMove(self.placeLatePose1Joint)
        self.robotMove(self.placeLatePose2Joint)


if __name__ == "__main__":
    rospy.init_node("robot_vision_wrapper")

    visionServerIP = rospy.get_param("robot_vision_wrapper/vision_server_ip")
    visionServerPort = rospy.get_param("robot_vision_wrapper/vision_server_port")
    isEyeInHand = rospy.get_param("robot_vision_wrapper/isEyeInHand")
    objectCompensationXYZ = rospy.get_param("robot_vision_wrapper/object_compensation_XYZ")
    useInputRPY = rospy.get_param("robot_vision_wrapper/use_input_RPY")
    objectRotateRPY = rospy.get_param("robot_vision_wrapper/object_rotate_RPY")
    objectPrePoseOffsetXYZRPY = rospy.get_param("robot_vision_wrapper/object_pre_pose_offset_XYZRPY")
    autoPlace = rospy.get_param("robot_vision_wrapper/auto_place")
    pickPrePose1Joint = rospy.get_param("robot_vision_wrapper/pick_pre_pose_1_joint")
    pickPrePose2Joint = rospy.get_param("robot_vision_wrapper/pick_pre_pose_2_joint")
    pickLatePose1Joint = rospy.get_param("robot_vision_wrapper/pick_late_pose_1_joint")
    pickLatePose2Joint = rospy.get_param("robot_vision_wrapper/pick_late_pose_2_joint")
    placeLatePose1Joint = rospy.get_param("robot_vision_wrapper/place_late_pose_1_joint")
    placeLatePose2Joint = rospy.get_param("robot_vision_wrapper/place_late_pose_2_joint")
    
    wrapper = RobotVisionWrapper()
    wrapper.initParam(isEyeInHand, objectCompensationXYZ, useInputRPY, objectRotateRPY, objectPrePoseOffsetXYZRPY, autoPlace,
                    pickPrePose1Joint, pickPrePose2Joint, pickLatePose1Joint, pickLatePose2Joint)

    wrapper.connectVision(visionServerIP, visionServerPort)
    wrapper.startVision()
    rospy.sleep(10)

    rospy.spin()

    wrapper.stopVision()
    wrapper.disconnectVision()
